package dd.lo.chapter22;

import dd.lo.util.ImageUtils;
import javafx.application.Application;
import javafx.application.Platform;
import javafx.concurrent.Task;
import javafx.scene.Scene;
import javafx.scene.control.Button;
import javafx.scene.image.Image;
import javafx.scene.image.ImageView;
import javafx.scene.layout.*;
import javafx.scene.paint.Color;
import javafx.scene.text.Text;
import javafx.stage.FileChooser;
import javafx.stage.Stage;
import org.opencv.core.*;
import org.opencv.face.LBPHFaceRecognizer;
import org.opencv.imgcodecs.Imgcodecs;
import org.opencv.imgproc.Imgproc;
import org.opencv.objdetect.CascadeClassifier;
import org.opencv.objdetect.Objdetect;

import java.io.ByteArrayInputStream;
import java.io.File;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/**
 * 人脸识别+人脸相似度计算（FaceRecognizer比较法）
 */
public class Example3 extends Application {
    private static final int DEFAULT_WIDTH = 1200;
    private static final int DEFAULT_HEIGHT = 900;

    private static final int FACE_WIDTH = 120;
    private static final int FACE_HEIGHT = 160;
    private static final double FACE_RATIO = (double) FACE_WIDTH / FACE_HEIGHT;

    private Application app;

    private boolean groupChooserClickable = false;

    private ImageView imageView1;

    private VBox subWin;

    private ImageView trainFacesImgView;

    private LBPHFaceRecognizer faceRecognizer;

    public static void main(String[] args) {
        System.loadLibrary(Core.NATIVE_LIBRARY_NAME);
        launch(args);
    }

    @Override
    public void start(Stage stage) {
        app = this;
        faceRecognizer = LBPHFaceRecognizer.create();
        trainFacesImgView = new ImageView();
        subWin = new VBox();
        subWin.getChildren().addAll(trainFacesImgView, new Text("line 1"), new Text("line 2"), new Text("line 3"));
        subWin.setBorder(new Border(new BorderStroke(Color.RED, BorderStrokeStyle.SOLID, CornerRadii.EMPTY, new BorderWidths(1))));
        Button faceChooserBtn = new Button("选择人脸图片");
        faceChooserBtn.setOnMouseClicked(e -> {
            FileChooser fileChooser = new FileChooser();
            fileChooser.setTitle("请选择目标人脸图像");
            fileChooser.setInitialDirectory(new File(System.getProperty("user.home") + "/Downloads")
            );
            fileChooser.getExtensionFilters().addAll(
//                    new FileChooser.ExtensionFilter("All Images", "*.*"),
                    new FileChooser.ExtensionFilter("JPEG", "*.jpeg"),
                    new FileChooser.ExtensionFilter("WEBP", "*.webp"),
                    new FileChooser.ExtensionFilter("PNG", "*.png"),
                    new FileChooser.ExtensionFilter("JPG", "*.jpg"),
                    new FileChooser.ExtensionFilter("GIF", "*.gif"),
                    new FileChooser.ExtensionFilter("BMP", "*.bmp")
            );
            List<File> imgFiles = fileChooser.showOpenMultipleDialog(stage);
            if (null == imgFiles || imgFiles.size() < 1) {
                System.out.println("用户取消选中人脸文件");
                return;
            }
            //读取多张训练用头像
            Mat[] trainFaces = new Mat[imgFiles.size()];
            for (int i = 0; i < imgFiles.size(); ++i) {
                trainFaces[i] = Imgcodecs.imread(imgFiles.get(i).getAbsolutePath(), Imgcodecs.IMREAD_UNCHANGED);
                if (trainFaces[i].empty()) {
                    System.out.printf("%s图片读取失败\n", imgFiles.get(i).getAbsolutePath());
                    throw new RuntimeException("图片读取失败");
                }
            }
            trainFaces(trainFaces);
            groupChooserClickable = true;
        });
        Button groupChooserBtn = new Button("选择合照图片");
        groupChooserBtn.setOnMouseClicked(e -> {
            if (!groupChooserClickable) return;
            FileChooser fileChooser = new FileChooser();
            fileChooser.setTitle("请选择含有人脸的合照");
            fileChooser.setInitialDirectory(new File(System.getProperty("user.home") + "/Downloads")
            );
            fileChooser.getExtensionFilters().addAll(
//                    new FileChooser.ExtensionFilter("All Images", "*.*"),
                    new FileChooser.ExtensionFilter("JPEG", "*.jpeg"),
                    new FileChooser.ExtensionFilter("WEBP", "*.webp"),
                    new FileChooser.ExtensionFilter("PNG", "*.png"),
                    new FileChooser.ExtensionFilter("JPG", "*.jpg"),
                    new FileChooser.ExtensionFilter("GIF", "*.gif"),
                    new FileChooser.ExtensionFilter("BMP", "*.bmp")
            );
            File imgFile = fileChooser.showOpenDialog(stage);
            if (null == imgFile) {
                System.out.println("用户取消选中文件");
                return;
            }
            //读取一张图片
            Mat srcImg = Imgcodecs.imread(imgFile.getAbsolutePath(), Imgcodecs.IMREAD_UNCHANGED);
            if (srcImg.empty()) {
                throw new RuntimeException("图片读取失败");
            }
            detectFace(srcImg);
        });
        FlowPane bottomControl = new FlowPane();
        bottomControl.getChildren().addAll(groupChooserBtn, faceChooserBtn);
        bottomControl.setPrefHeight(50);
        BorderPane root = new BorderPane();
        imageView1 = new ImageView();//imageView放到一个pane的中间
        imageView1.setFitWidth(DEFAULT_WIDTH);
        root.setCenter(imageView1);
        root.setBottom(bottomControl);
        root.setRight(subWin);
        Scene scene = new Scene(root, DEFAULT_WIDTH + FACE_WIDTH, DEFAULT_HEIGHT + 50);
        stage.setScene(scene);
        stage.setTitle("Example2-Face Detect and Compare");
        stage.setResizable(false);
        stage.show();
    }

    private void trainFaces(Mat[] facesToBeTrained) {
        Task<Void> task = new Task<Void>() {
            @Override
            protected Void call() {
                List<Mat> faces = new ArrayList<>(facesToBeTrained.length);
                int[] labels = new int[facesToBeTrained.length];
                Arrays.fill(labels, 1);
                List<Mat> displayFaces = new ArrayList<>(facesToBeTrained.length);
                for (Mat trainFace : facesToBeTrained) {
                    //调整图片大小以用于展示和训练
                    Size imgFitSize = new Size(FACE_WIDTH, FACE_HEIGHT);
                    Imgproc.resize(trainFace, trainFace, imgFitSize);
                    displayFaces.add(trainFace);
                    Mat grayImg = new Mat();
                    Imgproc.cvtColor(trainFace, grayImg, Imgproc.COLOR_BGR2GRAY);
                    faces.add(grayImg);
                }
                faceRecognizer.train(faces, new MatOfInt(labels));
                //将训练用得人脸展示出来
                Mat displayImg = new Mat();
                Core.vconcat(displayFaces, displayImg);
                MatOfByte buffer = new MatOfByte();
                Imgcodecs.imencode(".jpeg", displayImg, buffer);
                trainFacesImgView.setImage(new Image(new ByteArrayInputStream(buffer.toArray())));
                return null;
            }
        };
        new Thread(task).start();
    }

    private void detectFace(Mat srcImg) {
        Task<Void> task = new Task<Void>() {
            @Override
            protected Void call() {
                double ratio = (double) srcImg.width() / srcImg.height();
                Mat grayImg = new Mat();
                Imgproc.cvtColor(srcImg, grayImg, Imgproc.COLOR_BGR2GRAY);
                MatOfRect results = new MatOfRect();
                CascadeClassifier classifier = new CascadeClassifier( app.getClass().getResource("").getPath() + "haarcascade_frontalface_alt.xml");
                int minWidth = (int) Math.round((double) srcImg.height() * 0.02);
                int minHeight = (int) Math.round((double) srcImg.height() * 0.02);
                classifier.detectMultiScale(grayImg, results, 1.1, 3, Objdetect.CASCADE_DO_CANNY_PRUNING, new Size(minWidth, minHeight));
                Rect[] faceRoiArr = results.toArray();
                double[] compareResults = new double[faceRoiArr.length];
                for (int i = 0; i < faceRoiArr.length; ++i) {
                    Rect rect = faceRoiArr[i];
                    //画出识别结果
//                    Imgproc.rectangle(srcImg, rect, new Scalar(0, 255, 0), 3);
                    Rect faceRoi = new Rect((int) (rect.x - rect.width * 0.1), (int) (rect.y - rect.height * 0.3), (int) (rect.width * 1.2), (int) (rect.height * 1.6));
                    Mat faceImg = new Mat(grayImg, faceRoi);
                    Imgproc.resize(faceImg, faceImg, new Size(FACE_WIDTH, FACE_HEIGHT));
                    //利用faceRecognizer预测两张脸是同一张的自信度
                    int[] predictLabel = new int[1];
                    double[] predictLoss = new double[1];
                    faceRecognizer.predict(faceImg, predictLabel, predictLoss);
                    compareResults[i] = predictLoss[0];
                    Imgproc.rectangle(srcImg, faceRoi, new Scalar(0, 255, 0), 3);
                    Imgproc.putText(srcImg, i + "", faceRoi.tl(), Imgproc.FONT_HERSHEY_SIMPLEX, 0.6, new Scalar(0, 0, 255), 1, Imgproc.LINE_8);
                }
                //调整图片大小以用于展示
                ImageUtils.fitImgSize(srcImg, DEFAULT_WIDTH, DEFAULT_HEIGHT);
                //将最终结果展示出来
                MatOfByte buffer = new MatOfByte();
                Imgcodecs.imencode(".jpeg", srcImg, buffer);
                imageView1.setImage(new Image(new ByteArrayInputStream(buffer.toArray())));
                Platform.runLater(() -> {
                    subWin.getChildren().clear();
                    subWin.getChildren().add(trainFacesImgView);
                    subWin.getChildren().addAll(new Text("预测目标人脸偏差值"), new Text("（数值越低越可能是目标人物）"));
                    for (int i = 0; i < compareResults.length; ++i) {
                        subWin.getChildren().add(new Text(String.format("%d: %f", i, compareResults[i])));
                    }
                });
                return null;
            }
        };
        new Thread(task).start();
    }
}
